Plotting grids¶

We've covered plotting data sets in the Introduction to Data Analysis course but sometimes you want a raw view of an array.

The main tool to do this in Python is matplotlib's imshow. This is designed for displaying images (which are stored as numpy arrays) but it works on any 2D array:

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(constrained_layout=True)
ax.axis("off")

grid = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
ax.imshow(grid)

<matplotlib.image.AxesImage at 0x7fb6bc1a8910>

We can use this to display the results of our simulation, as it creates a 2D array:

with np.load("infection_simulation.npz") as f:
    cells = f["state"]
print(cells)

[[0 1 0 0 1]
 [0 1 1 0 1]
 [0 1 1 1 1]
 [0 1 1 1 1]
 [0 1 1 1 1]]

passing cells to imshow:

fig, ax = plt.subplots(constrained_layout=True)
ax.axis("off")

ax.imshow(cells)

<matplotlib.image.AxesImage at 0x7fb6832fb0d0>

creates a plot with the default colour scheme. Usually with cellular automata you want to assign specific colours to specific states. In our case we want the healthy cells to be green and the infected cells to be orange.

If we define a colour map with ListedColormap we can give a list of colour values corresponding to the cell values.

We also specify the vmin and vmax of the data to make sure that it aligns with our list of colours:

from matplotlib.colors import ListedColormap

fig, ax = plt.subplots(constrained_layout=True)
ax.axis("off")

cmap = ListedColormap(["#029e73", "#de8f05"])

ax.imshow(
    cells,
    cmap=cmap,
    vmin=0,
    vmax=1,
)

<matplotlib.image.AxesImage at 0x7fb6832cb5b0>

We can interpret this in the same way as the print(cells) output. The first row is the initial state, then each row down the image is another time step.

As ever with a plot, we can save it with fig.savefig:

fig.savefig("cells.svg")

Exercise¶

Plot the summary grid for the one-dimensional cellular automaton simulation you ran from the last chapter.

Previous | Next

Numpy and Simulations

Plotting grids¶

Exercise¶